import os
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
import cv2
import numpy as np
import glob
from para import L,R
import math
# 找棋盘格角点
# 阈值
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
#棋盘格模板规格
w = 11   #内角点个数，内角点是和其他格子连着的点
h = 8

# 世界坐标系中的棋盘格点,例如(0,0,0), (1,0,0), (2,0,0) ....,(8,5,0)，去掉Z坐标，记为二维矩阵
objp = np.zeros((w*h,3), np.float32)
objp[:,:2] = np.mgrid[0:w,0:h].T.reshape(-1,2)
objp=objp*19.3#mm
#print("objp",objp)
# 储存棋盘格角点的世界坐标和图像坐标对
objpoints = [] # 在世界坐标系中的三维点
limgpoints = [] # 在图像平面的二维点
rimgpoints = []

#images = glob.glob('img414_1/*.jpg')+glob.glob('img414_2/*.jpg')
limages = glob.glob('img/L/*.jpg')
rimages = glob.glob('img/R/*.jpg')
def corn(fname):
    img = cv2.imread(fname)
    gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    ret, corners = cv2.findChessboardCorners(gray, (w,h),None)
    if ret == True:
        cv2.cornerSubPix(gray,corners,(11,11),(-1,-1),criteria)
        return corners
    else:
        return []
def corner(img):#找角点
    gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    ret, corners = cv2.findChessboardCorners(gray, (w,h),None)
    if ret == True:
        cv2.cornerSubPix(gray,corners,(11,11),(-1,-1),criteria)
        return corners
    else:
        return []
for i in range(len(limages)):
    lfname=r'img/L/L'+str(i)+'.jpg'
    rfname=r'img/R/R'+str(i)+'.jpg'
    print(lfname,rfname)
    l=corn(lfname)
    r=corn(rfname)
    if len(l)!=0 and len(r)!=0:
        objpoints.append(objp)
        limgpoints.append(l)
        rimgpoints.append(r)
    else:
        print("fail")
ret, L.imtx, L.dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, limgpoints, (1920,1080), None, None)

print(L.dist)

ret, R.imtx, R.dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, rimgpoints, (1920,1080), None, None)

print(R.dist)

#标定、去畸变
ret, L.imtx,L.dist,R.imtx,R.dist,Re,T,E,F = cv2.stereoCalibrate(objpoints, limgpoints,rimgpoints,L.imtx,L.dist,R.imtx,R.dist, (1920,1080) )

print(L.imtx)
print(R.imtx)
# mtx：内参数矩阵
# dist：畸变系数
# rvecs：旋转向量 （外参数）
# tvecs ：平移向量 （外参数）
print (("ret:"),ret)
print (("Re:\n"),Re)
print (("T:\n"),T)
print (("E:\n"),E)
print (("F:\n"),F)
print(np.linalg.norm(T))

# 光轴矫正
R1, R2, P1, P2, Q, roi_left, roi_right = cv2.stereoRectify(L.imtx, L.dist, R.imtx,R.dist , (1920,1080), Re, T, flags=0, alpha=0)#cv2.CALIB_ZERO_DISPARITY
print("P1\n",P1)
print("P2\n",P2)
print("Q\n",Q)
print("roi_left\n",roi_left)
print("roi_right\n",roi_right)
map1x, map1y = cv2.initUndistortRectifyMap(L.imtx, L.dist, R1, P1, (1920, 1080), cv2.CV_32FC1)
map2x, map2y = cv2.initUndistortRectifyMap(R.imtx, R.dist, R2, P2, (1920, 1080), cv2.CV_32FC1)

cap = cv2.VideoCapture(0,cv2.CAP_MSMF)#打开内置摄像机
cap.set(cv2.CAP_PROP_FRAME_HEIGHT,1080)
cap.set(cv2.CAP_PROP_FRAME_WIDTH,3840)
while cap.isOpened():#当摄像头打开时
    ret,frame=cap.read()#读取当前摄像头画面
    l=frame[:,0:1920]
    r=frame[:,1920:3840]
    #l=cv2.imread(r"img/L/L1.jpg")
    #r=cv2.imread(r"img/R/R1.jpg")

    rectifyed_l=cv2.remap(l,map1x,map1y,cv2.INTER_AREA)
    rectifyed_r=cv2.remap(r,map2x,map2y,cv2.INTER_AREA)

    #cv2.imwrite(r"img/rectifyed_l.jpg",rectifyed_l)
    #cv2.imwrite(r"img/rectifyed_r.jpg",rectifyed_r)

    l_corners=corner(rectifyed_l)
    #cv2.drawChessboardCorners(rectifyed_l, (w,h), l_corners, True)
    r_corners=corner(rectifyed_r)
    #cv2.drawChessboardCorners(rectifyed_r, (w,h), r_corners, True)
    if len(l_corners)==0 or len(r_corners)==0:
        continue
    pos=[]
    for i in range(w*h):
        a=(l_corners[i,0,0]-P1[0,2],l_corners[i,0,1])#转换为图像坐标系
        b=(r_corners[i,0,0]-P2[0,2],r_corners[i,0,1])
        #print("y:",a[1]-b[1])
        fl=P1[0,0]
        fr=P2[0,0]
        z=120/(a[0]/fl-b[0]/fr)#基线长度120mm,摄像机坐标系原点在左侧相机的镜头(小孔成像模型中的小孔)
        x=z*a[0]/fl
        y=z*a[1]/fr

        pos.append((x,y,z))
        x=round(x,1)
        y=round(y,1)
        z=round(z,1)
        cv2.putText(rectifyed_l,str(round(x-pos[0][0],1))+' '+str(round(y-pos[0][1],1))+' '+str(round(z-pos[0][2],1)),\
        (int(l_corners[i,0,0]),int(l_corners[i,0,1])),cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)#把在摄像机坐标系中的坐标标注在图上
        #cv2.putText(rectifyed_l,str(round(a[1]-b[1],1)),(int(l_corners[i,0,0]),int(l_corners[i,0,1])),cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)#把在摄像机坐标系中的坐标标注在图上
    #print(math.sqrt((pos[0][0]-pos[-1][0])**2+(pos[0][1]-pos[-1][1])**2+(pos[0][2]-pos[-1][2])**2))
    id=-1
    dis=math.sqrt((pos[0][0]-pos[id][0])**2+(pos[0][1]-pos[id][1])**2+(pos[0][2]-pos[id][2])**2)
    print("dis:",dis)
    #print(pos[0][0]-pos[5][0],pos[0][1]-pos[5][1],pos[0][2]-pos[5][2])
    cv2.imshow('rectifyed_l',rectifyed_l)
    cv2.waitKey(1)
